





<!DOCTYPE html>
<html class="writer-html5" lang="zh-CN" >
<head>
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>Use Tensorize to Leverage Hardware Intrinsics &mdash; tvm 0.8.dev1982 文档</title>
  

  
  <link rel="stylesheet" href="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/css/bootstrap.min.css" integrity="sha384-Gn5384xqQ1aoWXA+058RXPxPg6fy4IWvTNh0E263XmFcJlSAwiGgFAW/dAiS6JXm" crossorigin="anonymous">
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/gallery.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/css/tlcpack_theme.css" type="text/css" />

  
  
    <link rel="shortcut icon" href="../../_static/tvm-logo-square.png"/>
  

  
  
  
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../" src="../../_static/documentation_options.js"></script>
        <script data-url_root="../../" id="documentation_options" src="../../_static/documentation_options.js"></script>
        <script src="../../_static/jquery.js"></script>
        <script src="../../_static/underscore.js"></script>
        <script src="../../_static/doctools.js"></script>
        <script src="../../_static/translations.js"></script>
    
    <script type="text/javascript" src="../../_static/js/theme.js"></script>

    
    <script type="text/javascript" src="../../_static/js/tlcpack_theme.js"></script>
    <link rel="index" title="索引" href="../../genindex.html" />
    <link rel="search" title="搜索" href="../../search.html" />
    <link rel="next" title="Compute and Reduce with Tuple Inputs" href="tuple_inputs.html" />
    <link rel="prev" title="External Tensor Functions" href="extern_op.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    
<header class="header">
    <div class="innercontainer">
      <div class="headerInner d-flex justify-content-between align-items-center">
          <div class="headerLogo">
               <a href="https://tvm.apache.org/"><img src=https://tvm.apache.org/assets/images/logo.svg alt="logo"></a>
          </div>

          <div id="headMenu" class="headerNav">
            <button type="button" id="closeHeadMenu" class="navCloseBtn"><img src="../../_static/img/close-icon.svg" alt="Close"></button>
             <ul class="nav">
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/community>Community</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/download>Download</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/vta>VTA</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/blog>Blog</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/docs>Docs</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvmconf.org>Conference</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://github.com/apache/tvm/>Github</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvmchinese.github.io/declaration_zh_CN.html>About-Translators</a>
                </li>
             </ul>
               <div class="responsivetlcdropdown">
                 <button type="button" class="btn-link">
                   ASF
                 </button>
                 <ul>
                     <li>
                       <a href=https://apache.org/>Apache Homepage</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/licenses/>License</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/sponsorship.html>Sponsorship</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/security/>Security</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/thanks.html>Thanks</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/events/current-event>Events</a>
                     </li>
                     <li>
                       <a href=https://www.zhihu.com/column/c_1429578595417563136>Zhihu</a>
                     </li>
                 </ul>
               </div>
          </div>
            <div class="responsiveMenuIcon">
              <button type="button" id="menuBtn" class="btn-menu"><img src="../../_static/img/menu-icon.svg" alt="Menu Icon"></button>
            </div>

            <div class="tlcDropdown">
              <div class="dropdown">
                <button type="button" class="btn-link dropdown-toggle" data-toggle="dropdown" aria-haspopup="true" aria-expanded="false">
                  ASF
                </button>
                <div class="dropdown-menu dropdown-menu-right">
                  <ul>
                     <li>
                       <a href=https://apache.org/>Apache Homepage</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/licenses/>License</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/sponsorship.html>Sponsorship</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/security/>Security</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/thanks.html>Thanks</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/events/current-event>Events</a>
                     </li>
                     <li>
                       <a href=https://www.zhihu.com/column/c_1429578595417563136>Zhihu</a>
                     </li>
                  </ul>
                </div>
              </div>
          </div>
       </div>
    </div>
 </header>
 
    <nav data-toggle="wy-nav-shift" class="wy-nav-side fixed">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../index.html">
          

          
            
            <img src="../../_static/tvm-logo-small.png" class="logo" alt="Logo"/>
          
          </a>

          
            
            
                <div class="version">
                  0.8.dev1982
                </div>
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption" role="heading"><span class="caption-text">如何开始</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../install/index.html">安装 TVM</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../contribute/index.html">贡献者指南</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">用户引导</span></p>
<ul class="current">
<li class="toctree-l1"><a class="reference internal" href="../../tutorial/index.html">User Tutorial</a></li>
<li class="toctree-l1 current"><a class="reference internal" href="../index.html">How To Guides</a><ul class="current">
<li class="toctree-l2"><a class="reference internal" href="../compile_models/index.html">编译深度学习模型</a></li>
<li class="toctree-l2"><a class="reference internal" href="../deploy/index.html">TVM 部署模型和集成</a></li>
<li class="toctree-l2"><a class="reference internal" href="../work_with_relay/index.html">Work With Relay</a></li>
<li class="toctree-l2 current"><a class="reference internal" href="index.html">Work With Tensor Expression and Schedules</a><ul class="current">
<li class="toctree-l3"><a class="reference internal" href="schedule_primitives.html">TVM中的调度原语</a></li>
<li class="toctree-l3"><a class="reference internal" href="reduction.html">Reduction</a></li>
<li class="toctree-l3"><a class="reference internal" href="intrin_math.html">内部函数和数学函数</a></li>
<li class="toctree-l3"><a class="reference internal" href="scan.html">Scan and Recurrent Kernel</a></li>
<li class="toctree-l3"><a class="reference internal" href="extern_op.html">External Tensor Functions</a></li>
<li class="toctree-l3 current"><a class="current reference internal" href="#">Use Tensorize to Leverage Hardware Intrinsics</a><ul>
<li class="toctree-l4"><a class="reference internal" href="#define-matrix-multiplication">Define Matrix Multiplication</a></li>
<li class="toctree-l4"><a class="reference internal" href="#schedule-the-matmul">Schedule the Matmul</a></li>
<li class="toctree-l4"><a class="reference internal" href="#define-gemv-tensorization-intrinsic">Define GEMV Tensorization Intrinsic</a></li>
<li class="toctree-l4"><a class="reference internal" href="#reduce-update-for-tensorize">Reduce-update for Tensorize</a></li>
<li class="toctree-l4"><a class="reference internal" href="#summary">总结</a></li>
</ul>
</li>
<li class="toctree-l3"><a class="reference internal" href="tuple_inputs.html">Compute and Reduce with Tuple Inputs</a></li>
<li class="toctree-l3"><a class="reference internal" href="tedd.html">Use Tensor Expression Debug Display (TEDD) for Visualization</a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="../optimize_operators/index.html">优化张量算子</a></li>
<li class="toctree-l2"><a class="reference internal" href="../tune_with_autotvm/index.html">Auto-Tune with Templates and AutoTVM</a></li>
<li class="toctree-l2"><a class="reference internal" href="../tune_with_autoscheduler/index.html">Use AutoScheduler for Template-Free Scheduling</a></li>
<li class="toctree-l2"><a class="reference internal" href="../work_with_microtvm/index.html">Work With microTVM</a></li>
<li class="toctree-l2"><a class="reference internal" href="../extend_tvm/index.html">Extend TVM</a></li>
<li class="toctree-l2"><a class="reference internal" href="../profile/index.html">Profile Models</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../errors.html">Handle TVM Errors</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../faq.html">常见提问</a></li>
</ul>
</li>
</ul>
<p class="caption" role="heading"><span class="caption-text">开发者引导</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../dev/tutorial/index.html">Developer Tutorial</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../dev/how_to/how_to.html">开发者指南</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">架构指南</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../arch/index.html">Design and Architecture</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">主题引导</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../topic/microtvm/index.html">microTVM：裸机使用TVM</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../topic/vta/index.html">VTA: Versatile Tensor Accelerator</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">参考指南</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../reference/langref/index.html">语言参考</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/api/python/index.html">Python API</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/api/links.html">Other APIs</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/publications.html">Publications</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../genindex.html">索引</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">
      
      <nav class="wy-nav-top" aria-label="top navigation" data-toggle="wy-nav-top">
        
            <div class="togglemenu">

            </div>
            <div class="nav-content">
              <!-- tvm -->
              Table of content
            </div>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        

          




















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../index.html">Docs</a> <span class="br-arrow">></span></li>
        
          <li><a href="../index.html">How To Guides</a> <span class="br-arrow">></span></li>
        
          <li><a href="index.html">Work With Tensor Expression and Schedules</a> <span class="br-arrow">></span></li>
        
      <li>Use Tensorize to Leverage Hardware Intrinsics</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
            
            <a href="../../_sources/how_to/work_with_schedules/tensorize.rst.txt" rel="nofollow"> <img src="../../_static//img/source.svg" alt="viewsource"/></a>
          
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <div class="sphx-glr-download-link-note admonition note">
<p class="admonition-title">注解</p>
<p>点击 <a class="reference internal" href="#sphx-glr-download-how-to-work-with-schedules-tensorize-py"><span class="std std-ref">此处</span></a> 获取完整示例代码</p>
</div>
<div class="sphx-glr-example-title section" id="use-tensorize-to-leverage-hardware-intrinsics">
<span id="tutorials-tensorize"></span><span id="sphx-glr-how-to-work-with-schedules-tensorize-py"></span><h1>Use Tensorize to Leverage Hardware Intrinsics<a class="headerlink" href="#use-tensorize-to-leverage-hardware-intrinsics" title="永久链接至标题">¶</a></h1>
<p><strong>作者</strong>: <a class="reference external" href="https://github.com/yzhliu">Yizhi Liu</a></p>
<p>This is an introduction material on how to perform tensorization in TVM.</p>
<p>By using schedule primitive <code class="code docutils literal notranslate"><span class="pre">tensorize</span></code>,
people can replace a unit of computation with the corresponding intrinsics,
making it easy to leverage handcrafted micro-kernels,
as well as extend TVM to support new hardware architectures.</p>
<p>The purpose of this tutorial is to show the functionality
and usage of tensorize instead of providing an efficient solution.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">absolute_import</span><span class="p">,</span> <span class="n">print_function</span>

<span class="kn">import</span> <span class="nn">tvm</span>
<span class="kn">from</span> <span class="nn">tvm</span> <span class="k">import</span> <span class="n">te</span>
<span class="kn">import</span> <span class="nn">tvm.testing</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
</pre></div>
</div>
<div class="section" id="define-matrix-multiplication">
<h2>Define Matrix Multiplication<a class="headerlink" href="#define-matrix-multiplication" title="永久链接至标题">¶</a></h2>
<p>Take matrix multiplication as our example.
Matmul first multiply the corresponding elements between two matrix,
then accumulate across a certain axis.
The following lines describe the computation <code class="code docutils literal notranslate"><span class="pre">A</span> <span class="pre">*</span> <span class="pre">B^T</span></code> in TVM.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">N</span><span class="p">,</span> <span class="n">M</span><span class="p">,</span> <span class="n">L</span> <span class="o">=</span> <span class="mi">1024</span><span class="p">,</span> <span class="mi">512</span><span class="p">,</span> <span class="mi">64</span>
<span class="n">A</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">N</span><span class="p">,</span> <span class="n">L</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;A&quot;</span><span class="p">)</span>
<span class="n">B</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">M</span><span class="p">,</span> <span class="n">L</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;B&quot;</span><span class="p">)</span>
<span class="n">k</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">reduce_axis</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">L</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;k&quot;</span><span class="p">)</span>
<span class="n">C</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">compute</span><span class="p">((</span><span class="n">N</span><span class="p">,</span> <span class="n">M</span><span class="p">),</span> <span class="k">lambda</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">:</span> <span class="n">te</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">A</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">k</span><span class="p">]</span> <span class="o">*</span> <span class="n">B</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">k</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="n">k</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;C&quot;</span><span class="p">)</span>
<span class="n">s</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">create_schedule</span><span class="p">(</span><span class="n">C</span><span class="o">.</span><span class="n">op</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">lower</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">C</span><span class="p">],</span> <span class="n">simple_mode</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span>
</pre></div>
</div>
<p class="sphx-glr-script-out">输出:</p>
<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>primfn(A_1: handle, B_1: handle, C_1: handle) -&gt; ()
  attr = {&quot;from_legacy_te_schedule&quot;: True, &quot;global_symbol&quot;: &quot;main&quot;, &quot;tir.noalias&quot;: True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 64], []),
             B: Buffer(B_2: Pointer(float32), float32, [512, 64], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i: int32, 0, 1024) {
    for (j: int32, 0, 512) {
      C_2[((i*512) + j)] = 0f32
      for (k: int32, 0, 64) {
        C_2[((i*512) + j)] = ((float32*)C_2[((i*512) + j)] + ((float32*)A_2[((i*64) + k)]*(float32*)B_2[((j*64) + k)]))
      }
    }
  }
}
</pre></div>
</div>
</div>
<div class="section" id="schedule-the-matmul">
<h2>Schedule the Matmul<a class="headerlink" href="#schedule-the-matmul" title="永久链接至标题">¶</a></h2>
<p>Now, suppose we have an accelerator that supports
matrix-vector multiplication (GEMV) as a hardware primitive,
which can take arbitrary size of reduce axis,
but another axis needs to be no larger than 16.
Thus we break down the matmul loops to make the innermost loops a (16x64) GEMV.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">factor</span> <span class="o">=</span> <span class="mi">16</span>
<span class="n">x</span><span class="p">,</span> <span class="n">y</span> <span class="o">=</span> <span class="n">C</span><span class="o">.</span><span class="n">op</span><span class="o">.</span><span class="n">axis</span>
<span class="p">(</span><span class="n">z</span><span class="p">,)</span> <span class="o">=</span> <span class="n">C</span><span class="o">.</span><span class="n">op</span><span class="o">.</span><span class="n">reduce_axis</span>
<span class="n">yo</span><span class="p">,</span> <span class="n">yi</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">C</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">y</span><span class="p">,</span> <span class="n">factor</span><span class="o">=</span><span class="n">factor</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">C</span><span class="p">]</span><span class="o">.</span><span class="n">reorder</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">yo</span><span class="p">,</span> <span class="n">yi</span><span class="p">,</span> <span class="n">z</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">lower</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">C</span><span class="p">],</span> <span class="n">simple_mode</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span>
</pre></div>
</div>
<p class="sphx-glr-script-out">输出:</p>
<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>primfn(A_1: handle, B_1: handle, C_1: handle) -&gt; ()
  attr = {&quot;from_legacy_te_schedule&quot;: True, &quot;global_symbol&quot;: &quot;main&quot;, &quot;tir.noalias&quot;: True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 64], []),
             B: Buffer(B_2: Pointer(float32), float32, [512, 64], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i: int32, 0, 1024) {
    for (j.outer: int32, 0, 32) {
      for (j.inner: int32, 0, 16) {
        C_2[(((i*512) + (j.outer*16)) + j.inner)] = 0f32
        for (k: int32, 0, 64) {
          C_2[(((i*512) + (j.outer*16)) + j.inner)] = ((float32*)C_2[(((i*512) + (j.outer*16)) + j.inner)] + ((float32*)A_2[((i*64) + k)]*(float32*)B_2[(((j.outer*1024) + (j.inner*64)) + k)]))
        }
      }
    }
  }
}
</pre></div>
</div>
<p>As showed in the IR printed above,
the inner loops <code class="code docutils literal notranslate"><span class="pre">j.inner</span></code> along with <code class="code docutils literal notranslate"><span class="pre">k</span></code> together form a computation of GEMV
- within the inner most two loops, the index <code class="code docutils literal notranslate"><span class="pre">i</span></code> is fixed,
the access to the matrix <code class="code docutils literal notranslate"><span class="pre">A</span></code> only varies by <code class="code docutils literal notranslate"><span class="pre">k</span></code>,
which makes the access pattern of <code class="code docutils literal notranslate"><span class="pre">A</span></code> a “vector”.
In order to leverage our hypothetical hardware’s GEMV instruction,
we can tensorize over <code class="code docutils literal notranslate"><span class="pre">j.inner</span></code>.</p>
</div>
<div class="section" id="define-gemv-tensorization-intrinsic">
<h2>Define GEMV Tensorization Intrinsic<a class="headerlink" href="#define-gemv-tensorization-intrinsic" title="永久链接至标题">¶</a></h2>
<p>Before scheduling the tensorization, we need to first define the intrinsic function for GEMV.
It includes two parts, the first is a compute definition of GEMV.
TVM uses it to match the computing pattern in the original Matmul schedule.
The second is to specify how to execute GEMV on the device,
which is done in <code class="code docutils literal notranslate"><span class="pre">intrin_func</span></code> below.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">intrin_gemv</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">l</span><span class="p">):</span>
    <span class="n">a</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">l</span><span class="p">,),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;a&quot;</span><span class="p">)</span>
    <span class="n">b</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">m</span><span class="p">,</span> <span class="n">l</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;b&quot;</span><span class="p">)</span>
    <span class="n">k</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">reduce_axis</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">l</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;k&quot;</span><span class="p">)</span>
    <span class="n">c</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">compute</span><span class="p">((</span><span class="n">m</span><span class="p">,),</span> <span class="k">lambda</span> <span class="n">i</span><span class="p">:</span> <span class="n">te</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">a</span><span class="p">[</span><span class="n">k</span><span class="p">]</span> <span class="o">*</span> <span class="n">b</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">k</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="n">k</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;c&quot;</span><span class="p">)</span>
    <span class="n">Ab</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span><span class="n">a</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">a</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;A&quot;</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">strides</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
    <span class="n">Bb</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span><span class="n">b</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;B&quot;</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">strides</span><span class="o">=</span><span class="p">[</span><span class="n">te</span><span class="o">.</span><span class="n">var</span><span class="p">(</span><span class="s2">&quot;s1&quot;</span><span class="p">),</span> <span class="mi">1</span><span class="p">])</span>
    <span class="n">Cb</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span><span class="n">c</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">c</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;C&quot;</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">strides</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>

    <span class="k">def</span> <span class="nf">intrin_func</span><span class="p">(</span><span class="n">ins</span><span class="p">,</span> <span class="n">outs</span><span class="p">):</span>
        <span class="n">ib</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">ir_builder</span><span class="o">.</span><span class="n">create</span><span class="p">()</span>
        <span class="n">aa</span><span class="p">,</span> <span class="n">bb</span> <span class="o">=</span> <span class="n">ins</span>
        <span class="n">cc</span> <span class="o">=</span> <span class="n">outs</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">ib</span><span class="o">.</span><span class="n">emit</span><span class="p">(</span>
            <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">call_extern</span><span class="p">(</span>
                <span class="s2">&quot;int32&quot;</span><span class="p">,</span>
                <span class="s2">&quot;gemv_update&quot;</span><span class="p">,</span>
                <span class="n">cc</span><span class="o">.</span><span class="n">access_ptr</span><span class="p">(</span><span class="s2">&quot;w&quot;</span><span class="p">),</span>
                <span class="n">aa</span><span class="o">.</span><span class="n">access_ptr</span><span class="p">(</span><span class="s2">&quot;r&quot;</span><span class="p">),</span>
                <span class="n">bb</span><span class="o">.</span><span class="n">access_ptr</span><span class="p">(</span><span class="s2">&quot;r&quot;</span><span class="p">),</span>
                <span class="n">m</span><span class="p">,</span>
                <span class="n">l</span><span class="p">,</span>
                <span class="n">bb</span><span class="o">.</span><span class="n">strides</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
            <span class="p">)</span>
        <span class="p">)</span>
        <span class="k">return</span> <span class="n">ib</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>

    <span class="k">return</span> <span class="n">te</span><span class="o">.</span><span class="n">decl_tensor_intrin</span><span class="p">(</span><span class="n">c</span><span class="o">.</span><span class="n">op</span><span class="p">,</span> <span class="n">intrin_func</span><span class="p">,</span> <span class="n">binds</span><span class="o">=</span><span class="p">{</span><span class="n">a</span><span class="p">:</span> <span class="n">Ab</span><span class="p">,</span> <span class="n">b</span><span class="p">:</span> <span class="n">Bb</span><span class="p">,</span> <span class="n">c</span><span class="p">:</span> <span class="n">Cb</span><span class="p">})</span>
</pre></div>
</div>
<p>Here <code class="code docutils literal notranslate"><span class="pre">te.decl_tensor_intrin</span></code> declares how to execute the computation <code class="code docutils literal notranslate"><span class="pre">c.op</span></code>.
Our implementation simply takes the inputs and outputs,
converts them to pointers and emit an external function call.
Note that tensorization requires user to specify <code class="code docutils literal notranslate"><span class="pre">offset_factor</span></code>,
with this information, TVM has knowledge of whether the data is aligned
between the start address of the original data structure
and the offset being passed to tensorize,
so that it has chance to optimize with vectorized loading.
We set the factor to 1 for simplification.</p>
<p>Buffers are also declared for inputs and outputs, though this is not required,
we benefit from the extra information provided by buffers. For example, we pass
<code class="code docutils literal notranslate"><span class="pre">bb.strides[0]</span></code> as an argument to the external function <code class="code docutils literal notranslate"><span class="pre">gemv_update</span></code>.
For now <code class="code docutils literal notranslate"><span class="pre">bb.strides[0]</span> <span class="pre">==</span> <span class="pre">l</span></code>,
but later we will see how they can differ with more complicated schedules.</p>
<p>Note that we use <code class="code docutils literal notranslate"><span class="pre">te.var(&quot;s1&quot;)</span></code> as the first stride dimension for <code class="code docutils literal notranslate"><span class="pre">B</span></code>.
If the strides can be inferred
- in this case, TVM knows tensor B is compact thus the strides are <code class="code docutils literal notranslate"><span class="pre">[L,</span> <span class="pre">1]</span></code> -
such placeholder can be put to let TVM automatically bind the inferred value for us.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">gemv</span> <span class="o">=</span> <span class="n">intrin_gemv</span><span class="p">(</span><span class="n">factor</span><span class="p">,</span> <span class="n">L</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">C</span><span class="p">]</span><span class="o">.</span><span class="n">tensorize</span><span class="p">(</span><span class="n">yi</span><span class="p">,</span> <span class="n">gemv</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">lower</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">C</span><span class="p">],</span> <span class="n">simple_mode</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span>
</pre></div>
</div>
<p class="sphx-glr-script-out">输出:</p>
<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>primfn(A_1: handle, B_1: handle, C_1: handle) -&gt; ()
  attr = {&quot;from_legacy_te_schedule&quot;: True, &quot;global_symbol&quot;: &quot;main&quot;, &quot;tir.noalias&quot;: True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 64], []),
             B: Buffer(B_2: Pointer(float32), float32, [512, 64], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i: int32, 0, 1024) {
    for (j.outer: int32, 0, 32) {
      @tir.call_extern(&quot;gemv_update&quot;, @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), C_2, ((i*512) + (j.outer*16)), 16, 2, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), A_2, (i*64), 64, 1, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), B_2, (j.outer*1024), 1024, 1, dtype=handle), 16, 64, 64, dtype=int32)
    }
  }
}
</pre></div>
</div>
<p>By tensorizing over <code class="code docutils literal notranslate"><span class="pre">yi</span></code>, the inner most two loops are
now replaced by the intrinsic function we defined before.
In order to build and run the module, let’s define the external function <code class="code docutils literal notranslate"><span class="pre">gemv_update</span></code>,
it is a naive implementation of GEMV, just for demonstration.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">gemv_impl</span><span class="p">():</span>
    <span class="n">cc_code</span> <span class="o">=</span> <span class="s2">&quot;&quot;&quot;</span>
<span class="s2">      extern &quot;C&quot; int gemv_update(float *cc, float *aa, float *bb, int m, int l, int stride) {</span>
<span class="s2">        for (int i = 0; i &lt; m; ++i) {</span>
<span class="s2">            for (int j = 0; j &lt; l; ++j) {</span>
<span class="s2">                cc[i] += aa[j] * bb[i * stride + j];</span>
<span class="s2">            }</span>
<span class="s2">        }</span>
<span class="s2">        return 0;</span>
<span class="s2">      }</span>
<span class="s2">    &quot;&quot;&quot;</span>
    <span class="kn">from</span> <span class="nn">tvm.contrib</span> <span class="k">import</span> <span class="n">utils</span><span class="p">,</span> <span class="n">clang</span>

    <span class="n">temp</span> <span class="o">=</span> <span class="n">utils</span><span class="o">.</span><span class="n">tempdir</span><span class="p">()</span>
    <span class="n">ll_path</span> <span class="o">=</span> <span class="n">temp</span><span class="o">.</span><span class="n">relpath</span><span class="p">(</span><span class="s2">&quot;temp.ll&quot;</span><span class="p">)</span>
    <span class="c1"># Create LLVM ir from c source code</span>
    <span class="n">ll_code</span> <span class="o">=</span> <span class="n">clang</span><span class="o">.</span><span class="n">create_llvm</span><span class="p">(</span><span class="n">cc_code</span><span class="p">,</span> <span class="n">output</span><span class="o">=</span><span class="n">ll_path</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">ll_code</span>
</pre></div>
</div>
<p>Now we leverage the pragma attribute <code class="code docutils literal notranslate"><span class="pre">import_llvm</span></code> to import llvm asm inline.
The importing needs to happen before the tensorized GEMV being executed.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">s</span><span class="p">[</span><span class="n">C</span><span class="p">]</span><span class="o">.</span><span class="n">pragma</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="s2">&quot;import_llvm&quot;</span><span class="p">,</span> <span class="n">gemv_impl</span><span class="p">())</span>
<span class="nb">print</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">lower</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">C</span><span class="p">],</span> <span class="n">simple_mode</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span>
</pre></div>
</div>
<p class="sphx-glr-script-out">输出:</p>
<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>primfn(A_1: handle, B_1: handle, C_1: handle) -&gt; ()
  attr = {&quot;from_legacy_te_schedule&quot;: True, &quot;global_symbol&quot;: &quot;main&quot;, &quot;tir.noalias&quot;: True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 64], []),
             B: Buffer(B_2: Pointer(float32), float32, [512, 64], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  attr [IterVar(i: int32, (nullptr), &quot;DataPar&quot;, &quot;&quot;)] &quot;pragma_import_llvm&quot; = &quot;; ModuleID = &#39;/tmp/tmpnqhn4b2z/input0.cc&#39;
source_filename = &quot;/tmp/tmpnqhn4b2z/input0.cc&quot;
target datalayout = &quot;e-m:e-i64:64-f80:128-n8:16:32:64-S128&quot;
target triple = &quot;x86_64-pc-linux-gnu&quot;

; Function Attrs: noinline nounwind optnone uwtable
define dso_local i32 @gemv_update(float*, float*, float*, i32, i32, i32) #0 {
  %7 = alloca float*, align 8
  %8 = alloca float*, align 8
  %9 = alloca float*, align 8
  %10 = alloca i32, align 4
  %11 = alloca i32, align 4
  %12 = alloca i32, align 4
  %13 = alloca i32, align 4
  %14 = alloca i32, align 4
  store float* %0, float** %7, align 8
  store float* %1, float** %8, align 8
  store float* %2, float** %9, align 8
  store i32 %3, i32* %10, align 4
  store i32 %4, i32* %11, align 4
  store i32 %5, i32* %12, align 4
  store i32 0, i32* %13, align 4
  br label %15

15:                                               ; preds = %50, %6
  %16 = load i32, i32* %13, align 4
  %17 = load i32, i32* %10, align 4
  %18 = icmp slt i32 %16, %17
  br i1 %18, label %19, label %53

19:                                               ; preds = %15
  store i32 0, i32* %14, align 4
  br label %20

20:                                               ; preds = %46, %19
  %21 = load i32, i32* %14, align 4
  %22 = load i32, i32* %11, align 4
  %23 = icmp slt i32 %21, %22
  br i1 %23, label %24, label %49

24:                                               ; preds = %20
  %25 = load float*, float** %8, align 8
  %26 = load i32, i32* %14, align 4
  %27 = sext i32 %26 to i64
  %28 = getelementptr inbounds float, float* %25, i64 %27
  %29 = load float, float* %28, align 4
  %30 = load float*, float** %9, align 8
  %31 = load i32, i32* %13, align 4
  %32 = load i32, i32* %12, align 4
  %33 = mul nsw i32 %31, %32
  %34 = load i32, i32* %14, align 4
  %35 = add nsw i32 %33, %34
  %36 = sext i32 %35 to i64
  %37 = getelementptr inbounds float, float* %30, i64 %36
  %38 = load float, float* %37, align 4
  %39 = fmul float %29, %38
  %40 = load float*, float** %7, align 8
  %41 = load i32, i32* %13, align 4
  %42 = sext i32 %41 to i64
  %43 = getelementptr inbounds float, float* %40, i64 %42
  %44 = load float, float* %43, align 4
  %45 = fadd float %44, %39
  store float %45, float* %43, align 4
  br label %46

46:                                               ; preds = %24
  %47 = load i32, i32* %14, align 4
  %48 = add nsw i32 %47, 1
  store i32 %48, i32* %14, align 4
  br label %20

49:                                               ; preds = %20
  br label %50

50:                                               ; preds = %49
  %51 = load i32, i32* %13, align 4
  %52 = add nsw i32 %51, 1
  store i32 %52, i32* %13, align 4
  br label %15

53:                                               ; preds = %15
  ret i32 0
}

attributes #0 = { noinline nounwind optnone uwtable &quot;correctly-rounded-divide-sqrt-fp-math&quot;=&quot;false&quot; &quot;disable-tail-calls&quot;=&quot;false&quot; &quot;less-precise-fpmad&quot;=&quot;false&quot; &quot;min-legal-vector-width&quot;=&quot;0&quot; &quot;no-frame-pointer-elim&quot;=&quot;true&quot; &quot;no-frame-pointer-elim-non-leaf&quot; &quot;no-infs-fp-math&quot;=&quot;false&quot; &quot;no-jump-tables&quot;=&quot;false&quot; &quot;no-nans-fp-math&quot;=&quot;false&quot; &quot;no-signed-zeros-fp-math&quot;=&quot;false&quot; &quot;no-trapping-math&quot;=&quot;false&quot; &quot;stack-protector-buffer-size&quot;=&quot;8&quot; &quot;target-cpu&quot;=&quot;x86-64&quot; &quot;target-features&quot;=&quot;+cx8,+fxsr,+mmx,+sse,+sse2,+x87&quot; &quot;unsafe-fp-math&quot;=&quot;false&quot; &quot;use-soft-float&quot;=&quot;false&quot; }

!llvm.module.flags = !{!0}
!llvm.ident = !{!1}

!0 = !{i32 1, !&quot;wchar_size&quot;, i32 4}
!1 = !{!&quot;clang version 9.0.0-2~ubuntu18.04.2 (tags/RELEASE_900/final)&quot;}
&quot;;
  for (i, 0, 1024) {
    for (j.outer: int32, 0, 32) {
      @tir.call_extern(&quot;gemv_update&quot;, @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), C_2, ((i*512) + (j.outer*16)), 16, 2, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), A_2, (i*64), 64, 1, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), B_2, (j.outer*1024), 1024, 1, dtype=handle), 16, 64, 64, dtype=int32)
    }
  }
}
</pre></div>
</div>
<p>Finally we compare the tensorize version with that <code class="code docutils literal notranslate"><span class="pre">numpy.dot</span></code> produces,
ensure our implementation is correct.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">func</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">build</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">C</span><span class="p">],</span> <span class="n">target</span><span class="o">=</span><span class="s2">&quot;llvm&quot;</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;gemv&quot;</span><span class="p">)</span>

<span class="kn">from</span> <span class="nn">tvm.topi.utils</span> <span class="k">import</span> <span class="n">get_const_tuple</span>

<span class="n">dtype</span> <span class="o">=</span> <span class="n">A</span><span class="o">.</span><span class="n">dtype</span>
<span class="n">dev</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">device</span><span class="p">(</span><span class="s2">&quot;cpu&quot;</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
<span class="n">a</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">size</span><span class="o">=</span><span class="n">get_const_tuple</span><span class="p">(</span><span class="n">A</span><span class="o">.</span><span class="n">shape</span><span class="p">))</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">dtype</span><span class="p">)</span>
<span class="n">b</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">size</span><span class="o">=</span><span class="n">get_const_tuple</span><span class="p">(</span><span class="n">B</span><span class="o">.</span><span class="n">shape</span><span class="p">))</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">dtype</span><span class="p">)</span>
<span class="n">c</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">nd</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">get_const_tuple</span><span class="p">(</span><span class="n">C</span><span class="o">.</span><span class="n">shape</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">),</span> <span class="n">dev</span><span class="p">)</span>
<span class="n">func</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">nd</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">dev</span><span class="p">),</span> <span class="n">tvm</span><span class="o">.</span><span class="n">nd</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">dev</span><span class="p">),</span> <span class="n">c</span><span class="p">)</span>
<span class="n">tvm</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">assert_allclose</span><span class="p">(</span><span class="n">c</span><span class="o">.</span><span class="n">numpy</span><span class="p">(),</span> <span class="n">np</span><span class="o">.</span><span class="n">dot</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">T</span><span class="p">),</span> <span class="n">rtol</span><span class="o">=</span><span class="mf">1e-3</span><span class="p">)</span>
</pre></div>
</div>
</div>
<div class="section" id="reduce-update-for-tensorize">
<h2>Reduce-update for Tensorize<a class="headerlink" href="#reduce-update-for-tensorize" title="永久链接至标题">¶</a></h2>
<p>So far you have learned the basic idea of tensorize,
now let’s move one step forward to a more complicated case.</p>
<p>Assume our accelerator could only multiply a vector by a square matrix,
in which the vector size needs to be no larger than 16.
Given such hardware constrain, now we need to split the reduce axis as following,</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">zo</span><span class="p">,</span> <span class="n">zi</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">C</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="n">factor</span><span class="o">=</span><span class="n">factor</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">C</span><span class="p">]</span><span class="o">.</span><span class="n">reorder</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">yo</span><span class="p">,</span> <span class="n">zo</span><span class="p">,</span> <span class="n">yi</span><span class="p">,</span> <span class="n">zi</span><span class="p">)</span>
</pre></div>
</div>
<p>However, since the tensorize intrinsic now only covers a part of the reduce axis,
instead of using one “body” function, TVM requires a <code class="code docutils literal notranslate"><span class="pre">reduce_reset</span></code> function,
which will be invoked before the reduce for-loop, and a <code class="code docutils literal notranslate"><span class="pre">reduce_update</span></code> function,
which defines the “update” computing strategy.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">gemv_impl</span><span class="p">():</span>
    <span class="n">cc_code</span> <span class="o">=</span> <span class="s2">&quot;&quot;&quot;</span>
<span class="s2">      extern &quot;C&quot; int gemv_update(float *cc, float *aa, float *bb, int m, int l, int stride) {</span>
<span class="s2">        for (int i = 0; i &lt; m; ++i) {</span>
<span class="s2">            for (int j = 0; j &lt; l; ++j) {</span>
<span class="s2">                cc[i] += aa[j] * bb[i * stride + j];</span>
<span class="s2">            }</span>
<span class="s2">        }</span>
<span class="s2">        return 0;</span>
<span class="s2">      }</span>
<span class="s2">      extern &quot;C&quot; int gemv_reset(float *cc, int m) {</span>
<span class="s2">        for (int i = 0; i &lt; m; ++i) {</span>
<span class="s2">            cc[i] = 0.0;</span>
<span class="s2">        }</span>
<span class="s2">        return 0;</span>
<span class="s2">      }</span>
<span class="s2">    &quot;&quot;&quot;</span>
    <span class="kn">from</span> <span class="nn">tvm.contrib</span> <span class="k">import</span> <span class="n">utils</span><span class="p">,</span> <span class="n">clang</span>

    <span class="n">temp</span> <span class="o">=</span> <span class="n">utils</span><span class="o">.</span><span class="n">tempdir</span><span class="p">()</span>
    <span class="n">ll_path</span> <span class="o">=</span> <span class="n">temp</span><span class="o">.</span><span class="n">relpath</span><span class="p">(</span><span class="s2">&quot;temp.ll&quot;</span><span class="p">)</span>
    <span class="c1"># Create LLVM ir from c source code</span>
    <span class="n">ll_code</span> <span class="o">=</span> <span class="n">clang</span><span class="o">.</span><span class="n">create_llvm</span><span class="p">(</span><span class="n">cc_code</span><span class="p">,</span> <span class="n">output</span><span class="o">=</span><span class="n">ll_path</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">ll_code</span>


<span class="k">def</span> <span class="nf">intrin_gemv</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">l</span><span class="p">):</span>
    <span class="n">a</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">l</span><span class="p">,),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;a&quot;</span><span class="p">)</span>
    <span class="n">b</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">m</span><span class="p">,</span> <span class="n">l</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;b&quot;</span><span class="p">)</span>
    <span class="n">k</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">reduce_axis</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">l</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;k&quot;</span><span class="p">)</span>
    <span class="n">c</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">compute</span><span class="p">((</span><span class="n">m</span><span class="p">,),</span> <span class="k">lambda</span> <span class="n">i</span><span class="p">:</span> <span class="n">te</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">a</span><span class="p">[</span><span class="n">k</span><span class="p">]</span> <span class="o">*</span> <span class="n">b</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">k</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="n">k</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;c&quot;</span><span class="p">)</span>
    <span class="n">Ab</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span><span class="n">a</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">a</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;A&quot;</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">strides</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
    <span class="n">Bb</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span><span class="n">b</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;B&quot;</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">strides</span><span class="o">=</span><span class="p">[</span><span class="n">te</span><span class="o">.</span><span class="n">var</span><span class="p">(</span><span class="s2">&quot;s1&quot;</span><span class="p">),</span> <span class="mi">1</span><span class="p">])</span>
    <span class="n">Cb</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span><span class="n">c</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">c</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;C&quot;</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">strides</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>

    <span class="k">def</span> <span class="nf">intrin_func</span><span class="p">(</span><span class="n">ins</span><span class="p">,</span> <span class="n">outs</span><span class="p">):</span>
        <span class="n">aa</span><span class="p">,</span> <span class="n">bb</span> <span class="o">=</span> <span class="n">ins</span>
        <span class="n">cc</span> <span class="o">=</span> <span class="n">outs</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>

        <span class="k">def</span> <span class="nf">_body</span><span class="p">():</span>
            <span class="n">ib</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">ir_builder</span><span class="o">.</span><span class="n">create</span><span class="p">()</span>
            <span class="n">ib</span><span class="o">.</span><span class="n">emit</span><span class="p">(</span>
                <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">call_extern</span><span class="p">(</span>
                    <span class="s2">&quot;int32&quot;</span><span class="p">,</span>
                    <span class="s2">&quot;gemv_update&quot;</span><span class="p">,</span>
                    <span class="n">cc</span><span class="o">.</span><span class="n">access_ptr</span><span class="p">(</span><span class="s2">&quot;w&quot;</span><span class="p">),</span>
                    <span class="n">aa</span><span class="o">.</span><span class="n">access_ptr</span><span class="p">(</span><span class="s2">&quot;r&quot;</span><span class="p">),</span>
                    <span class="n">bb</span><span class="o">.</span><span class="n">access_ptr</span><span class="p">(</span><span class="s2">&quot;r&quot;</span><span class="p">),</span>
                    <span class="n">m</span><span class="p">,</span>
                    <span class="n">l</span><span class="p">,</span>
                    <span class="n">bb</span><span class="o">.</span><span class="n">strides</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
                <span class="p">)</span>
            <span class="p">)</span>
            <span class="k">return</span> <span class="n">ib</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>

        <span class="k">def</span> <span class="nf">_reduce_reset</span><span class="p">():</span>
            <span class="n">ib</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">ir_builder</span><span class="o">.</span><span class="n">create</span><span class="p">()</span>
            <span class="n">ib</span><span class="o">.</span><span class="n">emit</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">call_extern</span><span class="p">(</span><span class="s2">&quot;int32&quot;</span><span class="p">,</span> <span class="s2">&quot;gemv_reset&quot;</span><span class="p">,</span> <span class="n">cc</span><span class="o">.</span><span class="n">access_ptr</span><span class="p">(</span><span class="s2">&quot;w&quot;</span><span class="p">),</span> <span class="n">m</span><span class="p">))</span>
            <span class="k">return</span> <span class="n">ib</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>

        <span class="k">def</span> <span class="nf">_reduce_update</span><span class="p">():</span>
            <span class="k">return</span> <span class="n">_body</span><span class="p">()</span>

        <span class="k">return</span> <span class="n">_body</span><span class="p">(),</span> <span class="n">_reduce_reset</span><span class="p">(),</span> <span class="n">_reduce_update</span><span class="p">()</span>

    <span class="k">return</span> <span class="n">te</span><span class="o">.</span><span class="n">decl_tensor_intrin</span><span class="p">(</span><span class="n">c</span><span class="o">.</span><span class="n">op</span><span class="p">,</span> <span class="n">intrin_func</span><span class="p">,</span> <span class="n">binds</span><span class="o">=</span><span class="p">{</span><span class="n">a</span><span class="p">:</span> <span class="n">Ab</span><span class="p">,</span> <span class="n">b</span><span class="p">:</span> <span class="n">Bb</span><span class="p">,</span> <span class="n">c</span><span class="p">:</span> <span class="n">Cb</span><span class="p">})</span>
</pre></div>
</div>
<p>Note that <code class="code docutils literal notranslate"><span class="pre">intrin_func</span></code> now returns a triplet:
<code class="code docutils literal notranslate"><span class="pre">(body,</span> <span class="pre">reduce_reset,</span> <span class="pre">reduce_update)</span></code>.
If tensorization includes all the reduce axes, function <code class="code docutils literal notranslate"><span class="pre">body()</span></code> will be invoked,
otherwise <code class="code docutils literal notranslate"><span class="pre">reduce_reset()</span></code> and <code class="code docutils literal notranslate"><span class="pre">reduce_update()</span></code> together will be used.
In our example <code class="code docutils literal notranslate"><span class="pre">body()</span></code> and <code class="code docutils literal notranslate"><span class="pre">reduce_update()</span></code>
share the same implementation,
while in other cases, hardware may have different instructions for these two functions.
Moreover, we can see now <code class="code docutils literal notranslate"><span class="pre">bb.strides[0]</span></code> is different from <code class="code docutils literal notranslate"><span class="pre">l</span></code>
due to the tiling.</p>
<p>Tensorize for squared GEMV, build and check the results,</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">gemv</span> <span class="o">=</span> <span class="n">intrin_gemv</span><span class="p">(</span><span class="n">factor</span><span class="p">,</span> <span class="n">factor</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">C</span><span class="p">]</span><span class="o">.</span><span class="n">tensorize</span><span class="p">(</span><span class="n">yi</span><span class="p">,</span> <span class="n">gemv</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">C</span><span class="p">]</span><span class="o">.</span><span class="n">pragma</span><span class="p">(</span><span class="n">yo</span><span class="p">,</span> <span class="s2">&quot;import_llvm&quot;</span><span class="p">,</span> <span class="n">gemv_impl</span><span class="p">())</span>

<span class="n">func</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">build</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">B</span><span class="p">,</span> <span class="n">C</span><span class="p">],</span> <span class="n">target</span><span class="o">=</span><span class="s2">&quot;llvm&quot;</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;gemv&quot;</span><span class="p">)</span>
<span class="n">a</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">size</span><span class="o">=</span><span class="n">get_const_tuple</span><span class="p">(</span><span class="n">A</span><span class="o">.</span><span class="n">shape</span><span class="p">))</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">dtype</span><span class="p">)</span>
<span class="n">b</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">size</span><span class="o">=</span><span class="n">get_const_tuple</span><span class="p">(</span><span class="n">B</span><span class="o">.</span><span class="n">shape</span><span class="p">))</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">dtype</span><span class="p">)</span>
<span class="n">c</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">nd</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">get_const_tuple</span><span class="p">(</span><span class="n">C</span><span class="o">.</span><span class="n">shape</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">),</span> <span class="n">dev</span><span class="p">)</span>
<span class="n">func</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">nd</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">dev</span><span class="p">),</span> <span class="n">tvm</span><span class="o">.</span><span class="n">nd</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="n">dev</span><span class="p">),</span> <span class="n">c</span><span class="p">)</span>
<span class="n">tvm</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">assert_allclose</span><span class="p">(</span><span class="n">c</span><span class="o">.</span><span class="n">numpy</span><span class="p">(),</span> <span class="n">np</span><span class="o">.</span><span class="n">dot</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="o">.</span><span class="n">T</span><span class="p">),</span> <span class="n">rtol</span><span class="o">=</span><span class="mf">1e-3</span><span class="p">)</span>
</pre></div>
</div>
</div>
<div class="section" id="summary">
<h2>总结<a class="headerlink" href="#summary" title="永久链接至标题">¶</a></h2>
<p>This tutorial demonstrates the usage of tensorize intrinsic in TVM.
Tensorize provides a way for users to get fully optimized schedule via micro-kernels.
For example, INT8 quantization on Intel CPUs uses tensorization
to invoke AVX instruction directly.
It also enables TVM to compile to ASICs -
checkout <a class="reference internal" href="../../topic/vta/index.html#vta-index"><span class="std std-ref">VTA: Versatile Tensor Accelerator</span></a> for details.
We also demonstrates how to use inline assembly importing,
which helps users inject asm easily into the schedule.</p>
<div class="sphx-glr-footer class sphx-glr-footer-example docutils container" id="sphx-glr-download-how-to-work-with-schedules-tensorize-py">
<div class="sphx-glr-download docutils container">
<p><a class="reference download internal" download="" href="../../_downloads/428c6201e29ce74e73c6b41eee589f62/tensorize.py"><code class="xref download docutils literal notranslate"><span class="pre">Python</span> <span class="pre">源码下载:</span> <span class="pre">tensorize.py</span></code></a></p>
</div>
<div class="sphx-glr-download docutils container">
<p><a class="reference download internal" download="" href="../../_downloads/3b5e41b16a898b72d18127ebe2182c66/tensorize.ipynb"><code class="xref download docutils literal notranslate"><span class="pre">Jupyter</span> <span class="pre">notebook</span> <span class="pre">下载:</span> <span class="pre">tensorize.ipynb</span></code></a></p>
</div>
</div>
<p class="sphx-glr-signature"><a class="reference external" href="https://sphinx-gallery.github.io">Gallery generated by Sphinx-Gallery</a></p>
</div>
</div>


           </div>
           
          </div>
          

<footer>

    <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
      
        <a href="tuple_inputs.html" class="btn btn-neutral float-right" title="Compute and Reduce with Tuple Inputs" accesskey="n" rel="next">下一个 <span class="fa fa-arrow-circle-right"></span></a>
      
      
        <a href="extern_op.html" class="btn btn-neutral float-left" title="External Tensor Functions" accesskey="p" rel="prev"><span class="fa fa-arrow-circle-left"></span> 上一个</a>
      
    </div>

<div id="button" class="backtop"><img src="../../_static//img/right.svg" alt="backtop"/> </div>
<section class="footerSec">
    <div class="footerHeader">
      <ul class="d-flex align-md-items-center justify-content-between flex-column flex-md-row">
        <li class="copywrite d-flex align-items-center">
          <h5 id="copy-right-info">© 2020 Apache Software Foundation | All right reserved</h5>
        </li>
      </ul>

    </div>

    <ul>
      <li class="footernote">Copyright © 2020 The Apache Software Foundation. Apache TVM, Apache, the Apache feather, and the Apache TVM project logo are either trademarks or registered trademarks of the Apache Software Foundation.</li>
    </ul>

</section>
</footer>
        </div>
      </div>

    </section>

  </div>
  

    <script src="https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.12.9/umd/popper.min.js" integrity="sha384-ApNbgh9B+Y1QKtv3Rn7W3mgPxhU9K/ScQsAP7hUibX39j7fakFPskvXusvfa0b4Q" crossorigin="anonymous"></script>
    <script src="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/js/bootstrap.min.js" integrity="sha384-JZR6Spejh4U02d8jOt6vLEHfe/JQGiRRSQQxSfFWpi1MquVdAyjUar5+76PVCmYl" crossorigin="anonymous"></script>

  </body>
  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
    <!-- Theme Analytics -->
    <script>
    (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){
      (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o),
      m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m)
    })(window,document,'script','https://www.google-analytics.com/analytics.js','ga');

    ga('create', 'UA-75982049-2', 'auto');
    ga('send', 'pageview');
    </script>

    
   

</body>
</html>